import csv
import random
from datetime import datetime, timedelta

# ========================
# 生成 companies.csv
# ========================
with open('../static/file_csv/companies.csv', 'w', encoding='utf-8', newline='') as f:
    writer = csv.writer(f)
    writer.writerow([':id', 'name', 'reg_no', 'reg_date', 'status', 'legal_rep',
                     'registered_capital', 'industry', 'risk_level', ':LABEL'])
    industries = ['科技', '金融', '制造', '贸易', '房地产']
    risk_levels = ['低', '中', '高']
    for i in range(1, 101):
        reg_no = f"91310100{str(i).zfill(9)}"
        reg_date = (datetime(2020, 1, 1) + timedelta(days=i * 3)).strftime('%Y-%m-%d')
        status = '注销' if i % 10 == 0 else '存续'
        legal_rep = f"法人{(i % 50) + 1}"
        capital = ((random.randint(1, 100) + 1) * 100000)
        industry = industries[i % 5]
        risk_level = risk_levels[(i + i % 7) % 3]
        writer.writerow([f"C{i:03d}", f"公司{i}", reg_no, reg_date, status,
                         legal_rep, capital, industry, risk_level, "Company"])

# ========================
# 生成 shareholders.csv
# ========================
with open('../static/file_csv/shareholders.csv', 'w', encoding='utf-8', newline='') as f:
    writer = csv.writer(f)
    writer.writerow([':id', 'name', 'type', 'id_no', 'nationality', 'is_controlled', ':LABEL'])
    nationalities = ['中国', '美国', '新加坡']
    for i in range(1, 301):
        if i <= 200:
            name = f"自然人股东{i}"
            type_ = "自然人"
            id_no = f"ID{str(i).zfill(9)}"
        else:
            name = f"法人股东{i - 200}"
            type_ = "法人"
            id_no = f"91310100{str(i - 200).zfill(9)}"
        nationality = nationalities[i % 3]
        is_controlled = 'true' if i % 15 == 0 else 'false'
        writer.writerow([f"S{i:03d}", name, type_, id_no, nationality, is_controlled, "Shareholder"])

# ========================
# 生成 seniors.csv
# ========================
positions = ['董事长', 'CEO', 'CFO', '董事', '监事', '总经理']
with open('../static/file_csv/seniors.csv', 'w', encoding='utf-8', newline='') as f:
    writer = csv.writer(f)
    writer.writerow([':id', 'name', 'position', 'id_no', 'birth_date', 'gender', ':LABEL'])
    for i in range(1, 201):
        birth_year = 1970 + (i % 40)
        birth_date = f"{birth_year}-01-01"
        gender = '男' if i % 2 == 0 else '女'
        writer.writerow([f"E{i:03d}", f"高管{i}", positions[(i + i % 11) % 6],
                         f"ID{str(i).zfill(9)}", birth_date, gender, "Senior"])

# ========================
# 生成关系文件（需先读取节点 ID）
# ========================
# 收集所有 ID
company_ids = [f"C{i:03d}" for i in range(1, 101)]
shareholder_ids = [f"S{i:03d}" for i in range(1, 301)]
senior_ids = [f"E{i:03d}" for i in range(1, 201)]

# HOLDS: 股东 → 公司 (约 600 条)
holds_rels = []
while len(holds_rels) < 600:
    s = random.choice(shareholder_ids)
    c = random.choice(company_ids)
    ratio = round(random.uniform(0.01, 0.99), 2)
    if (s, c) not in [(r[0], r[1]) for r in holds_rels]:  # 去重
        holds_rels.append((s, c, ratio))

with open('../static/file_csv/holds.csv', 'w', encoding='utf-8', newline='') as f:
    writer = csv.writer(f)
    writer.writerow([':START_ID', ':END_ID', ':TYPE', 'ratio'])
    for s, c, r in holds_rels:
        writer.writerow([s, c, 'HOLDS', r])

# SERVES_AS: 高管 → 公司 (约 300 条)
serves_rels = []
while len(serves_rels) < 300:
    e = random.choice(senior_ids)
    c = random.choice(company_ids)
    if (e, c) not in [(r[0], r[1]) for r in serves_rels]:
        serves_rels.append((e, c))

with open('../static/file_csv/serves_as.csv', 'w', encoding='utf-8', newline='') as f:
    writer = csv.writer(f)
    writer.writerow([':START_ID', ':END_ID', ':TYPE', 'role'])
    roles = ['董事', '监事', '高级管理人员', '法定代表人']
    for e, c in serves_rels:
        writer.writerow([e, c, 'SERVES_AS', random.choice(roles)])

# INVESTS_IN: 公司 → 公司 (约 200 条)
invests_rels = []
while len(invests_rels) < 200:
    c1 = random.choice(company_ids)
    c2 = random.choice(company_ids)
    if c1 != c2 and (c1, c2) not in [(r[0], r[1]) for r in invests_rels]:
        invests_rels.append((c1, c2))

with open('../static/file_csv/invests_in.csv', 'w', encoding='utf-8', newline='') as f:
    writer = csv.writer(f)
    writer.writerow([':START_ID', ':END_ID', ':TYPE', 'amount'])
    for c1, c2 in invests_rels:
        amount = (random.randint(100, 5000) + 100) * 10000
        writer.writerow([c1, c2, 'INVESTS_IN', amount])

# RELATED_TO: 股东 ↔ 股东 (约 50 条)
related_rels = []
while len(related_rels) < 50:
    s1 = random.choice(shareholder_ids)
    s2 = random.choice(shareholder_ids)
    if s1 != s2 and (s1, s2) not in [(r[0], r[1]) for r in related_rels] and (s2, s1) not in [(r[0], r[1]) for r in related_rels]:
        related_rels.append((s1, s2))

with open('../static/file_csv/related_to.csv', 'w', encoding='utf-8', newline='') as f:
    writer = csv.writer(f)
    writer.writerow([':START_ID', ':END_ID', ':TYPE', 'relation'])
    relations = ['配偶', '父子', '母子', '兄弟', '一致行动人']
    for s1, s2 in related_rels:
        writer.writerow([s1, s2, 'RELATED_TO', random.choice(relations)])

print("✅ 所有 CSV 文件已生成！")